from flask_wtf import FlaskForm
from flask_wtf.file import FileAllowed, FileField, FileSize, FileRequired
from wtforms.fields.simple import StringField
from wtforms.validators import DataRequired
from src.entity.upload_entity import ALLOW_IMG_EXTENSION, ALLOW_AUDIO_EXTENSION


class ChatSchema(FlaskForm):
    query = StringField('query', validators=[
        DataRequired("query参数不能为空")
    ])


class UploadSchemaFile(FlaskForm):
    file = FileField(
        "file",
        validators=[
            FileRequired("上传图片不能为空"),
            FileSize(max_size=15 * 1024 * 1024, message="图片大小不能超过15MB"),
            FileAllowed(ALLOW_IMG_EXTENSION, message=f"仅允许上传{ALLOW_IMG_EXTENSION}"),
        ]
    )
class UploadImageSchema(FlaskForm):
    url = StringField(
        "url",
        validators=[
            DataRequired("url参数不能为空")
        ]
    )
    query = StringField('query', validators=[
        DataRequired("query参数不能为空")
    ])


class UploadAudioSchema(FlaskForm):
    file = FileField(
        "file",
        validators=[
            FileRequired("上传音频不能为空"),
            FileSize(max_size=15 * 1024 * 1024, message="音频大小不能超过15MB"),
            FileAllowed(ALLOW_AUDIO_EXTENSION, message=f"仅允许上传{ALLOW_AUDIO_EXTENSION}")
        ])


class UploadTextAudioSchema(FlaskForm):
    query = StringField('query', validators=[
        DataRequired("query参数不能为空")
    ])
